package systemservice

import (
    "fmt"
    "github.com/go-gst/go-gst/gst"
    "github.com/go-gst/go-gst/gst/video"
    "image"
    "image/color"
    "image/png"
    "os"
    "sync"
    "strings"
)

// Stream 表示一个视频流
type Stream struct {
    pipeline        *gst.Pipeline
    rtmpSink        *gst.Element
    h264FileSink    *gst.Element
    h265FileSink    *gst.Element
    appsink         *gst.Element
    isRTMPStreaming bool
    isSavingH264    bool
    isSavingH265    bool
    isDecodingRGB   bool
    codec           string
}

// StreamManager 管理多路视频流
type StreamManager struct {
    streams map[string]*Stream
    mutex   sync.Mutex
    nextID  int
}

// NewStreamManager 创建一个新的 StreamManager 实例
func NewStreamManager() *StreamManager {
    return &StreamManager{
        streams: make(map[string]*Stream),
        nextID:  1,
    }
}

// OpenStream 打开一路视频流
func (sm *StreamManager) OpenStream(rtspURL string) (string, error) {
    sm.mutex.Lock()
    defer sm.mutex.Unlock()

    // 生成唯一 ID
    id := fmt.Sprintf("%d", sm.nextID)
    sm.nextID++

    // 初始化 GStreamer
    gst.Init(nil)

    // 构建 GStreamer 管道
    pipelineStr := fmt.Sprintf(
        "rtspsrc location=%s ! {CODEC_DECODE} ! tee name=t " +
            "t. ! queue ! h264parse ! filesink name=h264filesink " +
            "t. ! queue ! h265parse ! filesink name=h265filesink " +
            "t. ! queue ! flvmux ! rtmpsink name=rtmpsink " +
            "t. ! queue ! {CODEC_DECODER} ! videoconvert ! video/x-raw,format=RGB ! appsink name=appsink",
        rtspURL)

    // 自动检测 H.264 或 H.265
    codecDecode := "rtph264depay ! h264parse"
    codecDecoder := "avdec_h264"
    codec := "h264"
    // 这里可以添加更复杂的检测逻辑，例如尝试解析流的 SDP 信息来确定编码格式

    pipelineStr = strings.ReplaceAll(pipelineStr, "{CODEC_DECODE}", codecDecode)
    pipelineStr = strings.ReplaceAll(pipelineStr, "{CODEC_DECODER}", codecDecoder)

    pipeline, err := gst.NewPipelineFromString(pipelineStr)
    if err != nil {
        return "", err
    }

    rtmpSink := pipeline.GetByName("rtmpsink")
    h264FileSink := pipeline.GetByName("h264filesink")
    h265FileSink := pipeline.GetByName("h265filesink")
    appsink := pipeline.GetByName("appsink")

    // 启动管道
    _, err = pipeline.SetState(gst.StatePlaying)
    if err != nil {
        return "", err
    }

    sm.streams[id] = &Stream{
        pipeline:        pipeline,
        rtmpSink:        rtmpSink,
        h264FileSink:    h264FileSink,
        h265FileSink:    h265FileSink,
        appsink:         appsink,
        isRTMPStreaming: false,
        isSavingH264:    false,
        isSavingH265:    false,
        isDecodingRGB:   false,
        codec:           codec,
    }

    return id, nil
}

// CloseStream 关闭一路视频流
func (sm *StreamManager) CloseStream(id string) error {
    sm.mutex.Lock()
    defer sm.mutex.Unlock()

    stream, exists := sm.streams[id]
    if!exists {
        return fmt.Errorf("stream with ID %s does not exist", id)
    }

    // 停止管道
    _, err := stream.pipeline.SetState(gst.StateNull)
    if err != nil {
        return err
    }

    delete(sm.streams, id)
    return nil
}

// StartRTMPStreaming 开始向 RTMP 推流
func (sm *StreamManager) StartRTMPStreaming(id, rtmpURL string) error {
    sm.mutex.Lock()
    defer sm.mutex.Unlock()

    stream, exists := sm.streams[id]
    if!exists {
        return fmt.Errorf("stream with ID %s does not exist", id)
    }

    if stream.isRTMPStreaming {
        return fmt.Errorf("RTMP streaming for stream %s is already started", id)
    }

    stream.rtmpSink.SetProperty("location", rtmpURL)
    _, err := stream.pipeline.SetState(gst.StatePlaying)
    if err != nil {
        return err
    }

    stream.isRTMPStreaming = true
    return nil
}

// StopRTMPStreaming 停止向 RTMP 推流
func (sm *StreamManager) StopRTMPStreaming(id string) error {
    sm.mutex.Lock()
    defer sm.mutex.Unlock()

    stream, exists := sm.streams[id]
    if!exists {
        return fmt.Errorf("stream with ID %s does not exist", id)
    }

    if!stream.isRTMPStreaming {
        return fmt.Errorf("RTMP streaming for stream %s is already stopped", id)
    }

    _, err := stream.pipeline.SetState(gst.StatePaused)
    if err != nil {
        return err
    }

    stream.isRTMPStreaming = false
    return nil
}

// StartSavingH264 开启保存 H.264 文件功能
func (sm *StreamManager) StartSavingH264(id, filePath string) error {
    sm.mutex.Lock()
    defer sm.mutex.Unlock()

    stream, exists := sm.streams[id]
    if!exists {
        return fmt.Errorf("stream with ID %s does not exist", id)
    }

    if stream.isSavingH264 {
        return fmt.Errorf("Saving H.264 for stream %s is already started", id)
    }

    stream.h264FileSink.SetProperty("location", filePath)
    _, err := stream.h264FileSink.SetState(gst.StatePlaying)
    if err != nil {
        return err
    }

    stream.isSavingH264 = true
    return nil
}

// StopSavingH264 关闭保存 H.264 文件功能
func (sm *StreamManager) StopSavingH264(id string) error {
    sm.mutex.Lock()
    defer sm.mutex.Unlock()

    stream, exists := sm.streams[id]
    if!exists {
        return fmt.Errorf("stream with ID %s does not exist", id)
    }

    if!stream.isSavingH264 {
        return fmt.Errorf("Saving H.264 for stream %s is already stopped", id)
    }

    _, err := stream.h264FileSink.SetState(gst.StatePaused)
    if err != nil {
        return err
    }

    stream.isSavingH264 = false
    return nil
}

// StartSavingH265 开启保存 H.265 文件功能
func (sm *StreamManager) StartSavingH265(id, filePath string) error {
    sm.mutex.Lock()
    defer sm.mutex.Unlock()

    stream, exists := sm.streams[id]
    if!exists {
        return fmt.Errorf("stream with ID %s does not exist", id)
    }

    if stream.isSavingH265 {
        return fmt.Errorf("Saving H.265 for stream %s is already started", id)
    }

    stream.h265FileSink.SetProperty("location", filePath)
    _, err := stream.h265FileSink.SetState(gst.StatePlaying)
    if err != nil {
        return err
    }

    stream.isSavingH265 = true
    return nil
}

// StopSavingH265 关闭保存 H.265 文件功能
func (sm *StreamManager) StopSavingH265(id string) error {
    sm.mutex.Lock()
    defer sm.mutex.Unlock()

    stream, exists := sm.streams[id]
    if!exists {
        return fmt.Errorf("stream with ID %s does not exist", id)
    }

    if!stream.isSavingH265 {
        return fmt.Errorf("Saving H.265 for stream %s is already stopped", id)
    }

    _, err := stream.h265FileSink.SetState(gst.StatePaused)
    if err != nil {
        return err
    }

    stream.isSavingH265 = false
    return nil
}

// StartDecodingRGB 开启解码为 RGB 图像功能
func (sm *StreamManager) StartDecodingRGB(id string) error {
    sm.mutex.Lock()
    defer sm.mutex.Unlock()

    stream, exists := sm.streams[id]
    if!exists {
        return fmt.Errorf("stream with ID %s does not exist", id)
    }

    if stream.isDecodingRGB {
        return fmt.Errorf("Decoding to RGB for stream %s is already started", id)
    }

    stream.isDecodingRGB = true
    return nil
}

// StopDecodingRGB 关闭解码为 RGB 图像功能
func (sm *StreamManager) StopDecodingRGB(id string) error {
    sm.mutex.Lock()
    defer sm.mutex.Unlock()

    stream, exists := sm.streams[id]
    if!exists {
        return fmt.Errorf("stream with ID %s does not exist", id)
    }

    if!stream.isDecodingRGB {
        return fmt.Errorf("Decoding to RGB for stream %s is already stopped", id)
    }

    stream.isDecodingRGB = false
    return nil
}

// GetRGBFrame 获取一帧 RGB 图像
func (sm *StreamManager) GetRGBFrame(id string) (image.Image, error) {
    sm.mutex.Lock()
    defer sm.mutex.Unlock()

    stream, exists := sm.streams[id]
    if!exists {
        return nil, fmt.Errorf("stream with ID %s does not exist", id)
    }

    if!stream.isDecodingRGB {
        return nil, fmt.Errorf("Decoding to RGB for stream %s is not started", id)
    }

    sample := stream.appsink.PullSample()
    if sample == nil {
        return nil, fmt.Errorf("no sample available from appsink")
    }
    defer sample.Unref()

    buffer := sample.GetBuffer()
    if buffer == nil {
        return nil, fmt.Errorf("no buffer in sample")
    }
    defer buffer.Unref()

    caps := sample.GetCaps()
    if caps == nil {
        return nil, fmt.Errorf("no caps in sample")
    }
    defer caps.Unref()

    info := gstvideo.NewVideoInfo()
    if!info.FromCaps(caps) {
        return nil, fmt.Errorf("failed to get video info from caps")
    }

    width := info.GetWidth()
    height := info.GetHeight()

    data, err := buffer.Map(gst.MapRead)
    if err != nil {
        return nil, err
    }
    defer buffer.Unmap()

    img := image.NewRGBA(image.Rect(0, 0, width, height))
    for y := 0; y < height; y++ {
        for x := 0; x < width; x++ {
            offset := (y*width + x) * 3
            r := data.Data[offset]
            g := data.Data[offset+1]
            b := data.Data[offset+2]
            img.SetRGBA(x, y, color.RGBA{r, g, b, 255})
        }
    }

    return img, nil
}
    